import numpy as np

from core.optimizer import Optimizer


class MomentumSGD(Optimizer):
    def __init__(self, lr=0.01, momentum=0.9):
        """
        带动量的随机梯度下降（Stochastic Gradient Descent, SGD）优化器。

        Args:
            lr (float): 学习率，控制参数更新的步长。
            momentum (float): 动量系数，通常为 0.9。
        """
        super().__init__()
        self.lr = lr
        self.momentum = momentum
        self.vs = {}  # 速度（动量）

    def update_one(self, param):
        """
        更新单个参数。

        Args:
            param (Parameter): 待更新的参数。

        Returns:
            None
        """
        # 获取参数的唯一标识
        v_key = id(param)
        # 如果速度项不存在，初始化为零数组
        if v_key not in self.vs:
            self.vs[v_key] = np.zeros_like(param.data)
        # 获取当前速度
        v = self.vs[v_key]
        # 计算动量项
        v = v * self.momentum
        # 更新速度
        v = v - self.lr * param.grad.data
        # 更新参数值
        param.data = param.data + v

